#include <dlfcn.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>

#include "jvmti.h"
#include "qemu-profile.h"

/* Tooling interface variables, there is always going to be one tooling 
 * interface running. So it is ok to declare them global */
static struct JNIInvokeInterface_ _invokeInterface;
static void *_agentHandle = 0;
static JavaVM _vm  = {&_invokeInterface};
static jvmtiNativeInterface _jvmtiNativeInterface;
static jvmtiEnv _env = { &_jvmtiNativeInterface };
static jvmtiEventCallbacks _eventCallbacks;
static jvmtiCapabilities _capabilities;
static bool _compiledMethodLoadNotify;
static bool _compiledMethodUnloadNotify;
static jlong _classTag = 0;

typedef struct 
{
   void *startPC;
   char *nameStr;
} method_signature;

/* Dummy functions for the JavaVM struct */
static jint DestroyJavaVM(JavaVM *vm) { return 0; }
static jint AttachCurrentThread(JavaVM *vm, void **penv, void *args) { return 0; }
static jint DetachCurrentThread(JavaVM *vm) { return 0; }
static jint AttachCurrentThreadAsDaemon(JavaVM *vm, void **penv, void *args) { return 0; }

static jvmtiEnv *getEnv(void) { return &_env; } 
static jvmtiCapabilities *getCapabilities(void) { return &_capabilities; } 
static jvmtiEventCallbacks *getEventCallbacks(void) { return &_eventCallbacks; } 
   
static void setCompiledMethodLoadNotification(bool mode) { _compiledMethodLoadNotify = mode; } 
static void setCompiledMethodUnloadNotification(bool mode) { _compiledMethodUnloadNotify = mode; } 
static void setClassTag(jlong  tag) { _classTag = tag; } 
static jlong getClassTag(void) { return _classTag; } 

/* JNI function implementations */
static jint GetEnv(JavaVM *vm, void **penv, jint version)
{
   // TODO: check this version stuff...  I don't think it's quite right
   if (version <= JVMTI_VERSION)
      {
      *penv = getEnv();
      return JNI_OK;
      }
   else
      {
      *penv = NULL;
      return JNI_EVERSION;
      }
}

/* JVMTI function implementations */
static jvmtiError GetPotentialCapabilities(jvmtiEnv* env, jvmtiCapabilities* capabilities_ptr)
{
   if (capabilities_ptr == NULL)
      return JVMTI_ERROR_NULL_POINTER;

   memset(capabilities_ptr, 0, sizeof(jvmtiCapabilities));
   capabilities_ptr->can_generate_compiled_method_load_events = 1;

   return JVMTI_ERROR_NONE;
}

static jvmtiError AddCapabilities(jvmtiEnv* env, const jvmtiCapabilities* capabilities_ptr)
{
   jvmtiCapabilities tmpCapabilities;

   //fprintf(stderr,"capabilities asked for %X\t%X\t%X\t%X size:%d\n",
   //    *(unsigned int*)capabilities_ptr, ((unsigned int*)capabilities_ptr)[1], ((unsigned int*)capabilities_ptr)[2],
   //     ((unsigned int*)capabilities_ptr)[3], sizeof(tmpCapabilities)); 
   
   if (capabilities_ptr == NULL)
      return JVMTI_ERROR_NULL_POINTER;

   memset(&tmpCapabilities, 0, sizeof(tmpCapabilities));
   
   tmpCapabilities.can_tag_objects = capabilities_ptr->can_tag_objects;
   getCapabilities()->can_tag_objects |= capabilities_ptr->can_tag_objects;
   tmpCapabilities.can_get_source_file_name = capabilities_ptr->can_get_source_file_name;
   getCapabilities()->can_get_source_file_name |= capabilities_ptr->can_get_source_file_name;
   tmpCapabilities.can_get_line_numbers = capabilities_ptr->can_get_line_numbers;
   getCapabilities()->can_get_line_numbers |= capabilities_ptr->can_get_line_numbers;
   tmpCapabilities.can_generate_compiled_method_load_events  = capabilities_ptr->can_generate_compiled_method_load_events;
   getCapabilities()->can_generate_compiled_method_load_events |= capabilities_ptr->can_generate_compiled_method_load_events;
   
   
   if (memcmp(&tmpCapabilities, capabilities_ptr, sizeof(tmpCapabilities)) == 0)
      return JVMTI_ERROR_NONE;
   

   return JVMTI_ERROR_NOT_AVAILABLE;
}

static jvmtiError RelinquishCapabilities(jvmtiEnv* env,	const jvmtiCapabilities* capabilities_ptr)
{
   if (capabilities_ptr == NULL)
      return JVMTI_ERROR_NULL_POINTER;

   if (capabilities_ptr->can_generate_compiled_method_load_events)
      getCapabilities()->can_generate_compiled_method_load_events = 0;

   return JVMTI_ERROR_NONE;
}

static jvmtiError GetCapabilities(jvmtiEnv* env, jvmtiCapabilities* capabilities_ptr)
{
   if (capabilities_ptr == NULL)
      return JVMTI_ERROR_NULL_POINTER;

   memcpy(capabilities_ptr, getCapabilities(), sizeof(jvmtiCapabilities));

   return JVMTI_ERROR_NONE;
}

static jvmtiError SetEventNotificationMode(jvmtiEnv* env, jvmtiEventMode mode, jvmtiEvent event_type, jthread event_thread, ...)
{
   // we don't support thread management
   if (event_thread)
      return JVMTI_ERROR_INVALID_THREAD;

   if (event_type < JVMTI_MIN_EVENT_TYPE_VAL || event_type > JVMTI_MAX_EVENT_TYPE_VAL)
      return JVMTI_ERROR_INVALID_EVENT_TYPE;

   if (mode != JVMTI_ENABLE && mode != JVMTI_DISABLE)
      return JVMTI_ERROR_ILLEGAL_ARGUMENT;

   //fprintf (stderr,"set event notification mode %d\n",event_type);fflush (stderr);

   switch (event_type)
      {
      case JVMTI_EVENT_COMPILED_METHOD_LOAD:
         if (getCapabilities()->can_generate_compiled_method_load_events) 
            // capabilities enablement and event enablement can be in any order?
            {
            setCompiledMethodLoadNotification(mode == JVMTI_ENABLE);
            return JVMTI_ERROR_NONE;
            }
         break;

      case JVMTI_EVENT_COMPILED_METHOD_UNLOAD:
         if (getCapabilities()->can_generate_compiled_method_load_events)
            {
            setCompiledMethodUnloadNotification(mode == JVMTI_ENABLE);
            return JVMTI_ERROR_NONE;
            }
         break;

      case JVMTI_EVENT_DYNAMIC_CODE_GENERATED:
      case JVMTI_EVENT_VM_INIT:
      case JVMTI_EVENT_VM_DEATH:
      case JVMTI_EVENT_CLASS_PREPARE:
         return JVMTI_ERROR_NONE;

      default:
         break;
      }

   return JVMTI_ERROR_MUST_POSSESS_CAPABILITY;
}

static jvmtiError CreateRawMonitor(jvmtiEnv* env,
                                   const char* name,
                                   jrawMonitorID* monitor_ptr)
{
      if ( monitor_ptr == NULL )
         return JVMTI_ERROR_NULL_POINTER;
      *monitor_ptr = 0;    
      return JVMTI_ERROR_NONE;
}

static jvmtiError DestroyRawMonitor(jvmtiEnv* env,
                                    jrawMonitorID monitor_id)
{
      return JVMTI_ERROR_NONE;
}

static jvmtiError RawMonitorEnter(jvmtiEnv* env,
                                  jrawMonitorID monitor_id)
{
      return JVMTI_ERROR_NONE;
}

static jvmtiError RawMonitorExit(jvmtiEnv* env,
                                  jrawMonitorID monitor_id)
{
      return JVMTI_ERROR_NONE;
}

static jvmtiError SetEventCallbacks(jvmtiEnv* env, const jvmtiEventCallbacks* callbacks, jint size_of_callbacks)
{
   if (size_of_callbacks < 0)
      return JVMTI_ERROR_ILLEGAL_ARGUMENT;
   
   uint32_t size = sizeof(jvmtiEventCallbacks);
   if (callbacks != NULL)
      {
      jvmtiEventCallbacks *cachedCallbacks = getEventCallbacks();
      memset(cachedCallbacks, 0, size);
      size = (size_of_callbacks < size) ? size_of_callbacks : size;
      memcpy(cachedCallbacks, callbacks, size);
      }
         
   return JVMTI_ERROR_NONE;
}

static jvmtiError Deallocate(jvmtiEnv* env, unsigned char* mem)
{
   if (mem)
      free(mem);

   return JVMTI_ERROR_NONE;
}

static jvmtiError GetTag(jvmtiEnv* env, jobject object, jlong* tag_ptr)
{
     *tag_ptr = getClassTag();
     // always return 0 if you want to create a new class load entry for each trace with tprof
     // *tag_ptr = 0;
     return JVMTI_ERROR_NONE;
}

static jvmtiError SetTag(jvmtiEnv* env, jobject object, jlong tag)
{
     setClassTag(tag);
     return JVMTI_ERROR_NONE;
}

static jvmtiError GetSourceFileName(jvmtiEnv* env, jclass klass,  char** source_name_ptr)
{
   if (source_name_ptr)
      {
      char c[] = "Trace.java";
      *source_name_ptr = (char *)malloc(sizeof (c));
      strcpy(*source_name_ptr,c);
      }
     return JVMTI_ERROR_NONE;
}

static jvmtiError GetClassMethods(jvmtiEnv* env, jclass klass, jint* method_count_ptr, jmethodID** methods_ptr)
{
     *method_count_ptr = 1 ;
     *methods_ptr = (jmethodID* )malloc(sizeof(jmethodID*));
     **methods_ptr = (jmethodID)klass;
     return JVMTI_ERROR_NONE;
}

static jvmtiError GetMethodDeclaringClass(jvmtiEnv* env, jmethodID method, jclass* declaring_class_ptr)
{
   if (declaring_class_ptr == NULL)
      return JVMTI_ERROR_NULL_POINTER;

   *declaring_class_ptr = (jclass)method;
   return JVMTI_ERROR_NONE;
}

static jvmtiError GetClassSignature(jvmtiEnv* env, jclass klass, char** signature_ptr, char** generic_ptr)
{
   if (signature_ptr)
      {
      char c[] = "Trace";
      *signature_ptr = (char *)malloc(sizeof (c));
      strcpy(*signature_ptr,c);
      }

   if (generic_ptr)
      {
      *generic_ptr = (char *)malloc(1);
      **generic_ptr = 0;
      }

   return JVMTI_ERROR_NONE;
}

static char *getSignature(void* startPC)
{
   static char trace_name[100];
   sprintf(trace_name, "%s-@0x%lX", "-qemu-translation-block", (unsigned long )startPC);
   return trace_name;
}

static jvmtiError GetMethodName(jvmtiEnv* env, jmethodID method, char** name_ptr, char** signature_ptr, char** generic_ptr)
{
   /* Return the signature of the method given it starting pc*/
   char *signature = getSignature(method);

   if (!signature)
      return JVMTI_ERROR_INVALID_METHODID;

   if (name_ptr)
      {
      *name_ptr = (char *)malloc(strlen(signature)+1);
      strcpy(*name_ptr, signature);
      }

   if (signature_ptr)
      {
      *signature_ptr = (char *)malloc(1);
      **signature_ptr = 0;
      }

   if (generic_ptr)
      {
      *generic_ptr = (char *)malloc(1);
      **generic_ptr = 0;
      }

   return JVMTI_ERROR_NONE;
}

static jvmtiError GetJLocationFormat(jvmtiEnv* env, jvmtiJlocationFormat* format_ptr)
{
   if (!format_ptr)
      return JVMTI_ERROR_NULL_POINTER;

   *format_ptr = JVMTI_JLOCATION_OTHER;

   return JVMTI_ERROR_NONE;
}

/* static members of tooling interface */

void initJVMTI(const char *agentFileName, const char *agentOptions)
{
   if (!agentFileName || agentFileName[0] == 0 || _agentHandle)
      return;

   memset(&_invokeInterface, 0, sizeof(_invokeInterface));
   _invokeInterface.DestroyJavaVM = DestroyJavaVM;
   _invokeInterface.AttachCurrentThread = AttachCurrentThread;
   _invokeInterface.DetachCurrentThread = DetachCurrentThread;
   _invokeInterface.GetEnv = GetEnv;
   _invokeInterface.AttachCurrentThreadAsDaemon = AttachCurrentThreadAsDaemon;


   /* Initialize the jvmti native interface, the interface is used to */
   memset(&_jvmtiNativeInterface, 0, sizeof(_jvmtiNativeInterface));
   
   _jvmtiNativeInterface.GetPotentialCapabilities = GetPotentialCapabilities;
   _jvmtiNativeInterface.AddCapabilities = AddCapabilities;
   _jvmtiNativeInterface.RelinquishCapabilities = RelinquishCapabilities;
   _jvmtiNativeInterface.GetCapabilities = GetCapabilities;
   _jvmtiNativeInterface.SetEventNotificationMode = SetEventNotificationMode;
   _jvmtiNativeInterface.SetEventCallbacks = SetEventCallbacks;
   _jvmtiNativeInterface.Deallocate = Deallocate;
   _jvmtiNativeInterface.GetMethodDeclaringClass = GetMethodDeclaringClass;
   _jvmtiNativeInterface.GetClassSignature = GetClassSignature;
   _jvmtiNativeInterface.GetMethodName = GetMethodName;
   _jvmtiNativeInterface.GetJLocationFormat = GetJLocationFormat;
   _jvmtiNativeInterface.CreateRawMonitor = CreateRawMonitor;
   _jvmtiNativeInterface.DestroyRawMonitor = DestroyRawMonitor;
   _jvmtiNativeInterface.RawMonitorEnter = RawMonitorEnter;
   _jvmtiNativeInterface.RawMonitorExit = RawMonitorExit;
   _jvmtiNativeInterface.GetClassMethods = GetClassMethods;
   _jvmtiNativeInterface.GetTag = GetTag;
   _jvmtiNativeInterface.SetTag = SetTag;
   _jvmtiNativeInterface.GetSourceFileName = GetSourceFileName;

   /* Register the event call backs */
   memset(&_capabilities, 0, sizeof(_capabilities));
   memset(&_eventCallbacks, 0, sizeof(_eventCallbacks));

   _agentHandle = dlopen(agentFileName, RTLD_LOCAL | RTLD_LAZY);
   
   if (_agentHandle)
      {
      int (JNICALL * onLoad)(JavaVM *, char *, void *);
      *(void **)(&onLoad) = dlsym(_agentHandle, "Agent_OnLoad");

      if (onLoad)
         {
         onLoad(&_vm, (char*) agentOptions, NULL);
         }
      else
         {
         dlclose(_agentHandle);
         _agentHandle = 0;
         }
      }

   return;
}

void shutdownJVMTI(void)
{
   if (_agentHandle)
      {
      void (JNICALL * onUnload)(JavaVM *);
      *(void **)(&onUnload) = dlsym(_agentHandle, "Agent_OnUnload");      

      if (onUnload)
         onUnload(&_vm);

      dlclose(_agentHandle);
      _agentHandle = 0;
      }
   return;
}
 
void methodCompiled(void *startPC, int32_t codeSize, const char *signature, const char *hotness)
{
   if (codeSize == 0)
      return;

   if (!_compiledMethodLoadNotify ||
       !_eventCallbacks.CompiledMethodLoad)
      {
      fprintf(stderr, "event call backs not set up\n");
      return;
      }

   char *sigCopy = (char *)malloc(strlen(signature)+(hotness?strlen(hotness):0)+2);
   if (hotness)
      sprintf(sigCopy, "%s-%s", signature, hotness);
   else
      sprintf(sigCopy, "%s", signature);

   _eventCallbacks.CompiledMethodLoad(getEnv(), (jmethodID)startPC, codeSize, startPC, 0, 0, 0);
}

void methodUnloaded(void *startPC)
{
   if (_compiledMethodUnloadNotify && _eventCallbacks.CompiledMethodUnload != NULL)
      {
      _eventCallbacks.CompiledMethodUnload(getEnv(), (jmethodID)startPC,  startPC);
      }
    
   return;
}
